package com.tdk.mybatisplus.demo.common.util;

import java.nio.charset.StandardCharsets;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
import javax.crypto.Cipher;
import lombok.extern.slf4j.Slf4j;

/**
 * <p>
 * RSA非对称算法
 * </p>
 *
 * @author: taodingkai
 * @modified:
 * @since: 2022/5/26 17:01
 */
@Slf4j
public final class RsaUtil {

  private static final String RSA_PRIVATE_KEY_PREFIX = "-----BEGIN PRIVATE KEY-----";
  private static final String RSA_PRIVATE_KEY_SUFFIX = "-----END PRIVATE KEY-----";
  private static final String RSA_PUBLIC_KEY_PREFIX = "-----BEGIN PUBLIC KEY-----";
  private static final String RSA_PUBLIC_KEY_SUFFIX = "-----END PUBLIC KEY-----";
  private static final String BACKSLASH_N = "\\n";
  private static final String BACKSLASH_R = "\\r";
  /**
   * 算法
   */
  public static final String RSA_KEY_ALGORITHM = "RSA";
  /**
   * 类型 NODE-RSA 默认是这个
   */
  public static final String RSA_TYPE = "RSA/ECB/OAEPPadding";

  /**
   * 签名算法
   */
  public static final String SIGNATURE_ALGORITHM = "MD5withRSA";

  private RsaUtil() {
  }

  public static final Map<String, Object> KEY_MAP_CACHE = new HashMap<String, Object>(2){{
    putAll(initKey());
  }};

  /**
   * <p>
   * 私钥解密
   * </p>
   *
   * @param data
   * @return java.long.String
   * @author: taodingkai
   * @since: 2022/5/26 17:06
   */
  public static String decryptByPrivateKey(String data) throws Exception {
    byte[] byteArray  = Base64.getMimeDecoder().decode(StandardCharsets.UTF_8.encode(data).array());
    String privateKey = getPrivateKey();
    return new String(decryptByPrivateKey(byteArray, privateKey), StandardCharsets.UTF_8);
  }

  /**
   * <p>
   * 公钥解密
   * </p>
   *
   * @param data
   * @return java.long.String
   * @author: taodingkai
   * @since: 2022/5/26 17:06
   */
  public static String decryptByPublicKey(String data) throws Exception {
    byte[] byteArray = Base64.getMimeDecoder().decode(StandardCharsets.UTF_8.encode(data).array());
    String publicKey = getPublicKey();
    return new String(decryptByPublicKey(byteArray, publicKey), StandardCharsets.UTF_8);
  }

  /**
   * <p>
   * 私钥加密
   * </p>
   *
   * @param data
   * @return java.long.String
   * @author: taodingkai
   * @since: 2022/5/26 17:06
   */
  public static String encryptByPrivateKey(String data) throws Exception {
    byte[] byteArray  = data.getBytes(StandardCharsets.UTF_8);
    String privateKey = getPrivateKey();
    return Base64.getMimeEncoder().encodeToString(encryptByPrivateKey(byteArray, privateKey));
  }

  /**
   * <p>
   * 公钥加密
   * </p>
   *
   * @param data
   * @return java.long.String
   * @author: taodingkai
   * @since: 2022/5/26 17:06
   */
  public static String encryptByPublicKey(String data) throws Exception {
    byte[] byteArray  = data.getBytes(StandardCharsets.UTF_8);
    String publicKey = getPublicKey();
    return Base64.getMimeEncoder().encodeToString(encryptByPublicKey(byteArray, publicKey));
  }

  /**
   * <p>
   * 获取私钥
   * </p>
   *
   * @return java.lang.String
   * @author: taodingkai
   * @since: 2022/5/26 17:19
   */
  public static String getPrivateKey() {
    return getPrivateKey(KEY_MAP_CACHE);
  }

  /**
   * <p>
   * 获取公钥
   * </p>
   *
   * @return java.lang.String
   * @author: taodingkai
   * @since: 2022/5/26 17:20
   */
  public static String getPublicKey() {
    return getPublicKey(KEY_MAP_CACHE);
  }

  /**
   * <p>
   * 私钥签名
   * </p>
   *
   * @param data
   * @return java.lang.String
   * @author: taodingkai
   * @since: 2022/5/26 17:24
   */
  public static String signByPrivateKey(String data) throws Exception {
    byte[] byteArray = data.getBytes(StandardCharsets.UTF_8);
    return signByPrivateKey(byteArray, getPrivateKey());
  }

  /**
   * <p>
   * 公钥校验
   * </p>
   *
   * @param data
   * @param sign
   * @return boolean
   * @author: taodingkai
   * @since: 2022/5/26 17:24
   */
  public static boolean verifyByPublicKey(String data, String sign) throws Exception {
    byte[] byteArray = Base64.getMimeDecoder().decode(StandardCharsets.UTF_8.encode(data).array());
    String publicKey = getPublicKey();
    return verifyByPublicKey(byteArray, publicKey, sign);
  }

  /**
   * 私钥解密
   *
   * @param data 需解密的数据
   * @param privateKey 私钥
   */
  public static byte[] decryptByPrivateKey(byte[] data, String privateKey) throws Exception {
    byte[]              keyBytes            = Base64.getMimeDecoder().decode(privateKey);
    PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(keyBytes);
    KeyFactory          keyFactory          = KeyFactory.getInstance(RSA_KEY_ALGORITHM);
    Key                 key                 = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
    Cipher              cipher              = Cipher.getInstance(RSA_TYPE);
    cipher.init(Cipher.DECRYPT_MODE, key);
    return cipher.doFinal(data);
  }

  /**
   * 公钥解密
   *
   * @param data 需解密的数据
   * @param publicKey 公钥
   */
  public static byte[] decryptByPublicKey(byte[] data, String publicKey) throws Exception {
    byte[]             keyBytes           = Base64.getMimeDecoder().decode(publicKey);
    X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(keyBytes);
    KeyFactory         keyFactory         = KeyFactory.getInstance(RSA_KEY_ALGORITHM);
    Key                key                = keyFactory.generatePublic(x509EncodedKeySpec);
    Cipher             cipher             = Cipher.getInstance(RSA_TYPE);
    cipher.init(Cipher.DECRYPT_MODE, key);
    return cipher.doFinal(data);
  }

  /**
   * 私钥加密
   *
   * @param data 需加密的数据
   * @param privateKey 私钥
   */
  public static byte[] encryptByPrivateKey(byte[] data, String privateKey) throws Exception {
    byte[]              keyBytes            = Base64.getMimeDecoder().decode(privateKey);
    PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(keyBytes);
    KeyFactory          keyFactory          = KeyFactory.getInstance(RSA_KEY_ALGORITHM);
    Key                 key                 = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
    Cipher              cipher              = Cipher.getInstance(RSA_TYPE);
    cipher.init(Cipher.ENCRYPT_MODE, key);
    return cipher.doFinal(data);
  }

  /**
   * 公钥加密
   *
   * @param data 需加密的数据
   * @param publicKey 公钥
   */
  public static byte[] encryptByPublicKey(byte[] data, String publicKey) throws Exception {
    byte[]             keyBytes           = Base64.getMimeDecoder().decode(publicKey);
    X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(keyBytes);
    KeyFactory         keyFactory         = KeyFactory.getInstance(RSA_KEY_ALGORITHM);
    Key                key                = keyFactory.generatePublic(x509EncodedKeySpec);
    Cipher             cipher             = Cipher.getInstance(RSA_TYPE);
    cipher.init(Cipher.ENCRYPT_MODE, key);
    return cipher.doFinal(data);
  }

  /**
   * 获取私钥
   *
   * @param keyMap 密钥Map
   */
  public static String getPrivateKey(Map<String, Object> keyMap) {
    Key key = (Key) keyMap.get("PrivateKey");
    return Base64.getMimeEncoder().encodeToString(key.getEncoded())
        .replaceFirst(RSA_PRIVATE_KEY_PREFIX,"")
        .replaceFirst(RSA_PRIVATE_KEY_SUFFIX,"")
        .replaceAll(BACKSLASH_R,"")
        .replaceAll(BACKSLASH_N,"")
        ;
  }

  /**
   * 获取公钥
   *
   * @param keyMap 密钥Map
   */
  public static String getPublicKey(Map<String, Object> keyMap) {
    Key key = (Key) keyMap.get("PublicKey");
    return Base64.getMimeEncoder().encodeToString(key.getEncoded())
        .replaceFirst(RSA_PUBLIC_KEY_PREFIX,"")
        .replaceFirst(RSA_PUBLIC_KEY_SUFFIX,"")
        .replaceAll(BACKSLASH_R,"")
        .replaceAll(BACKSLASH_N,"")
        ;
  }

  /**
   * 初始化密钥
   */
  public static Map<String, Object> initKey() {
    Map<String, Object> keyMap     = new HashMap<>(2);
    try {
      KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(RSA_KEY_ALGORITHM);
      keyPairGenerator.initialize(1024);
      KeyPair             keyPair    = keyPairGenerator.generateKeyPair();
      RSAPublicKey        publicKey  = (RSAPublicKey) keyPair.getPublic();
      RSAPrivateKey       privateKey = (RSAPrivateKey) keyPair.getPrivate();
      keyMap.put("PublicKey", publicKey);
      keyMap.put("PrivateKey", privateKey);
      String publicKeyStr = getPublicKey(keyMap);
      log.info("RSA公钥:{}",publicKeyStr);
      String privateKeyStr = getPrivateKey(keyMap);
      log.info("RSA私钥:{}",privateKeyStr);
    }catch (Exception e){
      log.error("初始化密钥失败",e);
    }
    return keyMap;
  }

  /**
   * 私钥签名
   *
   * @param data 需签名的数据
   * @param privateKey 私钥
   */
  public static String signByPrivateKey(byte[] data, String privateKey) throws Exception {
    byte[]              keyBytes            = Base64.getMimeDecoder().decode(privateKey);
    PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(keyBytes);
    KeyFactory          keyFactory          = KeyFactory.getInstance(RSA_KEY_ALGORITHM);
    PrivateKey          key                 = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
    Signature           signature           = Signature.getInstance(SIGNATURE_ALGORITHM);
    signature.initSign(key);
    signature.update(data);
    return Base64.getMimeEncoder().encodeToString(signature.sign());
  }

  /**
   * 公钥校验
   *
   * @param data 需校验的数据
   * @param publicKey 公钥
   * @param sign 签名
   */
  public static boolean verifyByPublicKey(byte[] data, String publicKey, String sign) throws Exception {
    byte[]             keyBytes           = Base64.getMimeDecoder().decode(publicKey);
    X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(keyBytes);
    KeyFactory         keyFactory         = KeyFactory.getInstance(RSA_KEY_ALGORITHM);
    PublicKey          key                = keyFactory.generatePublic(x509EncodedKeySpec);
    Signature          signature          = Signature.getInstance(SIGNATURE_ALGORITHM);
    signature.initVerify(key);
    signature.update(data);
    return signature.verify(Base64.getMimeDecoder().decode(sign));
  }

  public static void main(String[] args) throws Exception {
    String key = encryptByPublicKey("测试");
    String data = decryptByPrivateKey(key);
    String signKey = signByPrivateKey("测试");
    boolean result = verifyByPublicKey("测试",signKey);
    System.out.println(result);
    System.out.println(data);
  }
}

